import numpy as np
from gplearn.genetic import SymbolicRegressor
from bishe_situations.utils import generate_data, plot_results, evaluate_model

# 定义目标函数
def target_function(X):
    return X**3 + 2*X**2 - X + 1

# 生成数据
X, y = generate_data(target_function, n_samples=1000, noise=0.1)
# 将y转换为1维数组
y = y.ravel()

# 创建符号回归模型
est_gp = SymbolicRegressor(population_size=500,
                          generations=20,
                          tournament_size=20,
                          stopping_criteria=0.01,
                          const_range=(-1, 1),
                          init_depth=(2, 6),
                          init_method='half and half',
                          function_set=('add', 'sub', 'mul', 'div', 'sin', 'cos'),
                          metric='mean absolute error',
                          parsimony_coefficient=0.001,
                          random_state=0)

# 训练模型
est_gp.fit(X, y)

# 预测
y_pred = est_gp.predict(X)

# 评估模型
mse, r2 = evaluate_model(y, y_pred)

# 打印找到的表达式
print("\n找到的表达式:")
print(est_gp._program)

# 绘制结果
plot_results(X, y, y_pred, "GPLearn符号回归结果") 